import sys
import cv2
import numpy as np
import json
from tqdm import *
from utils import *
from motionbert.mb_inferencer import motionbert_inferencer
from motionbert.lib.utils.vismo import render_and_save


import pdb

"""
本程序实现以下功能：
1. 用morionbert做3D人体姿态推断
2. 保存人体姿态推断数据至 datasets/emohugo_pose
3. 保存人体姿态推断的可视化结果至 datasets/emohugo_pose_visual
"""

"""
推断2d人体姿态
输入模态：2d人体姿态json
输出模态：3d人体姿态json文件和可视化视频
"""

def inference_3d_and_save(json_path, inferencer, reault_data_path, result_vis_path, block_lower, silent):
    assert os.path.exists(json_path)
    ensure_dir(os.path.abspath(os.path.dirname(reault_data_path)))

    if not silent:
        print("inferencing 3d")
    result_data = inferencer(json_path, block_lower)
    json_str = json.dumps(result_data.tolist(), indent=4)
    with open(reault_data_path,"w") as f:
        f.write(json_str)

    if not result_vis_path is None:
        if not silent:
            print("exporting 3d visual")
        video_path = emohugo_path_shift(result_vis_path, "pose_3d_visual", "video")
        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        render_and_save(result_data, result_vis_path, fps=fps, silent=silent)

if __name__ == "__main__":
    config_path = os.path.join(get_project_root(), "motionbert/configs/pose3d/MB_train_h36m.yaml")
    ckpt_path = os.path.join(get_project_root(),
                             "motionbert/checkpoint/pose3d/MB_train_h36m_dual_loss/finish_at_2024_4_14/best_epoch.bin")

    inferencer_3d = motionbert_inferencer(config_path, ckpt_path)

    # video_dataset_root = os.path.join(get_project_root(), "datasets/emohugo_video")
    
    json_path = os.path.join(get_project_root() , "datasets/emohugo_pose/happiness/000002-S001-0/000002-S001-0_2d.json")
    print("inferencing 3d from: ", json_path)

    reault_data_path = emohugo_path_shift(json_path, "pose_2d", "pose_3d")
    result_vis_path = emohugo_path_shift(json_path, "pose_2d", "pose_3d_visual")
    print("saving to: \n", reault_data_path, "\n", result_vis_path)
    inference_3d_and_save(json_path, inferencer_3d, reault_data_path, result_vis_path, block_lower=True, silent=False)
